{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G3MMAcssHTML"
      },
      "source": [
        "<link rel=\"stylesheet\" href=\"/site-assets/css/gemma.css\">\n",
        "<link rel=\"stylesheet\" href=\"https://fonts.googleapis.com/css2?family=Google+Symbols:opsz,wght,FILL,GRAD@20..48,100..700,0..1,-50..200\" />"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pYCKXrlPH0VW"
      },
      "source": [
        "# Using LIT with Gemma"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hy4PNm1FHMqB"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemma/docs/lit_gemma\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on Generative AI</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://codelabs.developers.google.com/codelabs/responsible-ai/lit-gemma\"><img src=\"https://www.tensorflow.org/images/codelabs_logo.png\" height=\"24\" width=\"48\" />Learn in Codelabs</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_JK95-LDQ8Ln"
      },
      "source": [
        "Generative AI products are relatively new and their behaviors can vary more than\n",
        "earlier forms of software. This makes it important to probe the machine learning\n",
        "models being used, examine examples of the model's behavior and investigate\n",
        "surprises.\n",
        "\n",
        "The Learning Interpretability Tool (LIT; [website][lit-web], [GitHub][lit-gh])\n",
        "is a platform for debugging and analyzing ML models to understand why and how\n",
        "they behave the way they do.\n",
        "\n",
        "Here, you'll learn how to setup LIT to get more out of Google's \n",
        "[Gemma model][gemma] by using the Sequence Salience module to analyze different\n",
        "prompt engineering approaches.\n",
        "\n",
        "[lit-web]: https://pair-code.github.io/lit\n",
        "[lit-gh]: https://github.com/PAIR-code/lit\n",
        "[gemma]: https://ai.google.dev/gemma\n",
        "[keras-nlp]: https://keras.io/keras_nlp/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W3DeT6ysMb7g"
      },
      "source": [
        "# Setting up LIT to Debug Gemma Prompts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-8h-_gLNP6fd"
      },
      "source": [
        "*Note: you may see some warnings of the form*\n",
        "\n",
        "```\n",
        "ERROR: pip's dependency resolver does not currently take into account all the \n",
        "packages that are installed. This behaviour is the source of the following \n",
        "dependency conflicts.\n",
        "bigframes 0.21.0 requires scikit-learn>=1.2.2, but you have scikit-learn 1.0.2 \n",
        "which is incompatible.\n",
        "google-colab 1.0.0 requires ipython==7.34.0, but you have ipython 8.14.0 \n",
        "which is incompatible.\n",
        "```\n",
        "\n",
        "*These are safe to ignore.*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUakOwZaJSa2"
      },
      "source": [
        "## Install LIT and Keras NLP\n",
        "\n",
        "This notebook uses the KerasNLP implementation of Gemma (more on how to \n",
        "configure this below). You will need a recent version of `keras` (3.0+) \n",
        "`keras-nlp` (0.12+) and `lit-nlp` (1.2+), and a Kaggle account to download the \n",
        "base model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D5ktQxc6Jbo5"
      },
      "outputs": [],
      "source": [
        "# Keras is included in Colab runtimes, but needs to be updated to to v3.0+.\n",
        "# LIT and Keras NLP are not icnldued by default and must be installed.\n",
        "# Running this cell may require you to restart your session to ensure the newer\n",
        "# packages are imported correctly.\n",
        "! pip install -q -U \"keras >= 3.0, <4.0\" \"keras-nlp >= 0.14\" \"lit-nlp >= 1.2\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VzFhJDayXroT"
      },
      "source": [
        "### Kaggle Access\n",
        "\n",
        "KerasNLP stores their pre-trained model weights on Kaggle. The\n",
        "[`kagglehub` package](https://github.com/Kaggle/kagglehub#authenticate) is used\n",
        "to autheticate with this service. Be sure to also accept the license agreement \n",
        "for [Gemma](https://www.kaggle.com/models/keras/gemma) from your Kaggle account.\n",
        "\n",
        "See the Appendix at the end for more information on how to set up a Kaggle\n",
        "account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yKw8gDsh_nVR"
      },
      "outputs": [],
      "source": [
        "import kagglehub\n",
        "\n",
        "kagglehub.login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f0-KhV-rShMa"
      },
      "source": [
        "## Configuring LIT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v12sUSdiNPje"
      },
      "source": [
        "LIT provides a function, `make_notebook_widget()` for configuring our prompt \n",
        "debugging tools in a notebook context. \n",
        "\n",
        "LIT provides a dataset of sample prompts that accompany the tutorial linked \n",
        "later in this document.\n",
        "\n",
        "See the comments below for configuring the widget to use different models and/or\n",
        "datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BdgJqj56QSOC"
      },
      "outputs": [],
      "source": [
        "from lit_nlp.examples.prompt_debugging import notebook as lit_pdbnb\n",
        "\n",
        "# The following function initializes a LIT Notebook Widget. It's configured by\n",
        "# two required positional arguments:\n",
        "#\n",
        "# * `datasets_config`: A list of strings containing the dataset names and\n",
        "#       paths to load from, as \"dataset:path\", where path can be a URL or a\n",
        "#       local file path. The example below uses a special value,\n",
        "#       `sample_prompts`, to load the example prompts provided in the LIT\n",
        "#       distribution; no other special values are supported.\n",
        "# * `models_config`: A list of strings containing the model names and paths to\n",
        "#       load from, as \"model:path\", where path can be a URL, a local file path,\n",
        "#       or the name of a preset for the configured deep learning framework.\n",
        "#\n",
        "# LIT supports salience computation for KerasNLP and Hugging Face Transformers\n",
        "# models running on TensorFlow or PyTorch. Note that all models passed to the\n",
        "# `models_config` parameter will be loaded using the same framework and runtime.\n",
        "# You can cofnigre these with the following keywork arguments.\n",
        "#\n",
        "# * `dl_framework`: Must be one of \"kerasnlp\" or \"transformers\".\n",
        "# * `dl_runtime`: Must be one of \"tensorflow\" or \"torch\".\n",
        "#\n",
        "# Changing the `dl_framework` value will affect the authentication method used\n",
        "# to access Gemma model weights.\n",
        "\n",
        "lit_widget = lit_pdbnb.make_notebook_widget(\n",
        "    ['sample_prompts'],\n",
        "    [\"gemma_2b_it:gemma_1.1_instruct_2b_en\"],\n",
        "    dl_framework=\"kerasnlp\",\n",
        "    dl_runtime=\"tensorflow\",\n",
        "    batch_size=1,\n",
        "    max_examples=5,\n",
        "    precision=\"bfloat16\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q6snqitV2wPX"
      },
      "source": [
        "Now you can render the UI in a Colab cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F9DZMtEfURob"
      },
      "outputs": [
        {
          "data": {
            "application/javascript": [
              "(async (port, path, width, height, cache, element) => {\n",
              "    if (!google.colab.kernel.accessAllowed && !cache) {\n",
              "      return;\n",
              "    }\n",
              "    element.appendChild(document.createTextNode(''));\n",
              "    const url = await google.colab.kernel.proxyPort(port, {cache});\n",
              "    const iframe = document.createElement('iframe');\n",
              "    iframe.src = new URL(path, url).toString();\n",
              "    iframe.height = height;\n",
              "    iframe.width = width;\n",
              "    iframe.style.border = 0;\n",
              "    iframe.allow = [\n",
              "        'accelerometer',\n",
              "        'autoplay',\n",
              "        'camera',\n",
              "        'clipboard-read',\n",
              "        'clipboard-write',\n",
              "        'gyroscope',\n",
              "        'magnetometer',\n",
              "        'microphone',\n",
              "        'serial',\n",
              "        'usb',\n",
              "        'xr-spatial-tracking',\n",
              "    ].join('; ');\n",
              "    element.appendChild(iframe);\n",
              "  })(36717, \"/?\", \"100%\", \"800\", false, window.element)"
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "lit_widget.render()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AawFyPUUS4cW"
      },
      "source": [
        "# Prompt Debugging with Sequence Salience\n",
        "\n",
        "Text-to-text large language models (LLMs), such as Gemma, take an input sequence \n",
        "in the form of [tokenized][tokenization] text and generate new tokens that are \n",
        "logical follow-ons or completions.\n",
        "\n",
        "[Salience methods][salience-explorable] allow you to inspect which parts of an \n",
        "input are important to the model for different parts of its generated output. \n",
        "LIT's [Sequence Salience module][lit-seq-sal] extends these methods to explain \n",
        "the importance of sequences at multiple levels of granularity: from tokens to \n",
        "words to sentences and beyond.\n",
        "\n",
        "You can use LIT in the cell above to play around with the Sequence Salience \n",
        "module on your own. For a more guided learning experience, you can follow long \n",
        "with the [_Prompt Debugging with Sequence Salience_ tutorial][seq-sal-tutorial]\n",
        "right in this Colab.\n",
        "\n",
        "For even more academic and techncial information on how Sequence Salience works,\n",
        "check out [our paper][seq-sal-paper].\n",
        "\n",
        "\n",
        "[lit-seq-sal]: https://pair-code.github.io/lit/documentation/components.html#sequence-salience\n",
        "[salience-explorable]: https://pair.withgoogle.com/explorables/saliency/\n",
        "[seq-sal-paper]: https://arxiv.org/abs/2404.07498\n",
        "[seq-sal-tutorial]: https://pair-code.github.io/lit/tutorials/sequence-salience/\n",
        "[tokenization]: https://arxiv.org/pdf/1808.06226.pdf\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IZWodzCgPM6k"
      },
      "source": [
        "# Appendix: Accessing Gemma on Kaggle Hub\n",
        "\n",
        "This notebook uses the KerasNLP implementation of Gemma in this document. \n",
        "KerasNLP stores their pre-trained model weights on Kaggle, and Gemma requires \n",
        "authentication and license acknowledgement to access those weights.\n",
        "\n",
        "The following instruction walk you through how to set up a Kaggle account and \n",
        "authenticate with Kaggle using the `kagglehub` package.\n",
        "\n",
        "1.  Create a Kaggle account if you don't have one\n",
        "    * Go to: https://www.kaggle.com/account/login?phase=startRegisterTab\n",
        "    * Use whichever registation method you prefer to set up your account.\n",
        "1. Request access to Gemma\n",
        "    * Make sure you're logged into Kaggle using the account above\n",
        "    * Go to the consent page: https://www.kaggle.com/models/google/gemma/license/consent\n",
        "    * Select the \"Verify via Kaggle Account\" option (the default selection) and click next\n",
        "    * Complete the consent form (first name and last name fields at the top)\n",
        "    * Acknowledge the policy using the checkboxes at the bottom\n",
        "    * Click the \"Accept\" button at the bottom to be granted access\n",
        "    * This should redirect you to the model page (https://www.kaggle.com/models/google/gemma)\n",
        "1. Create an API token\n",
        "    * Make sure you're logged into Kaggle using the account you created above\n",
        "    * Got to the Settings page: https://www.kaggle.com/settings\n",
        "    * Scroll down to the API section\n",
        "    * Use the \"Create New Token\" button to trigger token generation\n",
        "    * Use the on-screen menu to save the JSON file, named kaggle.json, the service generates to your machine\n",
        "    * The JSON file is an object with two properties, username and key, you'll need both to authenticate with their service later\n",
        "1. Use your API token credentials to authenticate with kagglehub in Colab\n",
        "    * Go to the LIT Sequence Saleince Colab: https://colab.sandbox.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb#scrollTo=yKw8gDsh_nVR\n",
        "    * Conntect to a GPU runtime\n",
        "    * For Gemma 2B you can use the free-tier T4 runtime\n",
        "    * For Gemma 7B you will need pre-paid Colab compute credits or a Colab Pro account to use a V100, L4, or A100 GPU\n",
        "    * Run the `kagglehub` code cell to display an HTML form that asks for your username and a token\n",
        "    * Copy the `username` field from the `kaggle.json` file you downloaded in the previous step and paste it into the `username` field in the form\n",
        "    * Copy the `key` field from the `kaggle.json` file you downloaded in the previous step and paste it into the `token` field in the form\n",
        "    * Click the login button to save these credentials in your runtime\n",
        "\n",
        "You will need to repeat the last step any time the Colab runtime is disconnected, as disconnection clears the cache the credentials are stored in."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "lit_gemma.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
